package com.enba.boot.generator.generator;

import com.baomidou.mybatisplus.generator.FastAutoGenerator;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.rules.DbColumnType;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import com.enba.boot.generator.properties.GeneratorProperties;
import java.sql.Types;
import java.util.Collections;

/** 代码生成器 Helper */
public class AutoGeneratorHelper {

  private final GeneratorProperties generatorProperties;

  public AutoGeneratorHelper(GeneratorProperties generatorProperties) {
    this.generatorProperties = generatorProperties;
  }

  /**
   * 代码生成器
   *
   * @param generatorProperties 代码生成器配置
   */
  public void generator(GeneratorProperties generatorProperties) {
    if (generatorProperties.isEnabled()) {
      FastAutoGenerator.create(
              generatorProperties.getUrl(),
              generatorProperties.getUsername(),
              generatorProperties.getPassword())
          .globalConfig(
              builder -> {
                builder
                    .author(generatorProperties.getAuthor()) // 设置作者
                    .enableSwagger() // 开启 swagger 模式
                    .disableOpenDir() // 关闭默认的打开目录
                    .outputDir(generatorProperties.getOutputDir()); // 指定输出目录
              })
          .dataSourceConfig(
              builder ->
                  builder.typeConvertHandler(
                      (globalConfig, typeRegistry, metaInfo) -> {
                        int typeCode = metaInfo.getJdbcType().TYPE_CODE;
                        if (typeCode == Types.SMALLINT) {
                          // 自定义类型转换
                          return DbColumnType.INTEGER;
                        }
                        return typeRegistry.getColumnType(metaInfo);
                      }))
          .packageConfig(
              builder ->
                  builder
                      .parent(generatorProperties.getParent()) // 设置父包名
                      .moduleName(generatorProperties.getModuleName()) // 设置父包模块名
                      .pathInfo(
                          Collections.singletonMap(
                              OutputFile.xml, generatorProperties.getXmlDir())) // 设置mapperXml生成路径
              )
          .strategyConfig(
              builder ->
                  builder
                      .addInclude(generatorProperties.getInclude()) // 设置需要生成的表名
                      .addTablePrefix(generatorProperties.getTablePrefix()) // 设置过滤表前缀
              )
          .templateEngine(new FreemarkerTemplateEngine()) // 使用Freemarker引擎模板
          .execute();
    }
  }

  /** 代码生成器 */
  public void generator() {
    generator(generatorProperties);
  }
}
